from re import T
import requests
import json
import uuid

def test_health_check():
    """测试健康检查端点"""
    try:
        response = requests.get("http://localhost:8013/health")
        print(f"健康检查状态码: {response.status_code}")
        print(f"健康检查响应: {response.json()}")
        print()
    except Exception as e:
        print(f"健康检查测试失败: {e}")
        print()

def test_sql_query():
    """测试SQL查询端点"""
    try:
        response = requests.post("http://localhost:8013/sql/query", json={
            "query": "描述一下Album表的结构",
            "userId": "user123",
            "conversationId": "conv456"
        })
        
        print(f"SQL查询状态码: {response.status_code}")
        print(f"SQL查询响应: {json.dumps(response.json(), ensure_ascii=False, indent=2)}")
        print()
        
        # 第二个查询测试
        response2 = requests.post("http://localhost:8013/sql/query", json={
            "query": "查询前5个艺术家的信息",
            "userId": "user123",
            "conversationId": "conv456"
        })
        
        print(f"第二个查询状态码: {response2.status_code}")
        print(f"第二个查询响应: {json.dumps(response2.json(), ensure_ascii=False, indent=2)}")
        print()
        
    except Exception as e:
        print(f"SQL查询测试失败: {e}")
        print()

def test_chat_completion():
    """测试聊天完成接口 - 非流式"""
    try:
        # 生成唯一的用户ID和会话ID
        user_id = f"test_user_{uuid.uuid4().hex[:8]}"
        conversation_id = f"test_conv_{uuid.uuid4().hex[:8]}"
        
        response = requests.post("http://localhost:8013/v1/chat/completions", json={
            "messages": [
                {"role": "user", "content": "列出数据库中的所有表格"}
            ],
            "stream": False,
            "userId": user_id,
            "conversationId": conversation_id
        })
        
        print(f"聊天完成状态码: {response.status_code}")
        print(f"聊天完成响应: {json.dumps(response.json(), ensure_ascii=False, indent=2)}")
        print()
        
    except Exception as e:
        
        print(f"聊天完成测试失败: {e}")
        print()

def test_chat_completion_stream():
    """测试聊天完成接口 - 流式"""
    try:
        # 生成唯一的用户ID和会话ID
        user_id = f"test_user_{uuid.uuid4().hex[:8]}"
        conversation_id = f"test_conv_{uuid.uuid4().hex[:8]}"
        
        response = requests.post("http://localhost:8013/v1/chat/completions", json={
            "messages": [
                {"role": "user", "content": "查询Artist表中的前3个艺术家"}
            ],
            "stream": True,
            "userId": user_id,
            "conversationId": conversation_id
        }, stream=True)
        
        print(f"流式聊天完成状态码: {response.status_code}")
        print("流式响应内容:")
        
        # 处理流式响应
        for line in response.iter_lines():
            if line:
                line_str = line.decode('utf-8')
                if line_str.startswith('data: '):
                    data_str = line_str[6:]  # 移除 'data: ' 前缀
                    if data_str.strip():
                        try:
                            data = json.loads(data_str)
                            if 'choices' in data and data['choices']:
                                delta = data['choices'][0].get('delta', {})
                                content = delta.get('content', '')
                                if content:
                                    print(content, end='', flush=True)
                        except json.JSONDecodeError:
                            continue
        print("\n")
        print()
        
    except Exception as e:
        print(f"流式聊天完成测试失败: {e}")
        print()

def test_chat_completion_with_context():
    """测试带上下文的聊天完成接口"""
    try:
        # 使用相同的会话ID来测试上下文保持
        user_id = "context_test_user"
        conversation_id = "context_test_conv"
        
        # 第一个问题
        response1 = requests.post("http://localhost:8013/v1/chat/completions", json={
            "messages": [
                {"role": "user", "content": "Album表有多少条记录？"}
            ],
            "stream": False,
            "userId": user_id,
            "conversationId": conversation_id
        })
        
        print(f"第一个问题状态码: {response1.status_code}")
        print(f"第一个问题响应: {json.dumps(response1.json(), ensure_ascii=False, indent=2)}")
        print()
        
        # 第二个问题（基于上下文）
        response2 = requests.post("http://localhost:8013/v1/chat/completions", json={
            "messages": [
                {"role": "user", "content": "Album表有多少条记录？"},
                {"role": "assistant", "content": response1.json().get('choices', [{}])[0].get('message', {}).get('content', '')},
                {"role": "user", "content": "那Artist表呢？"}
            ],
            "stream": False,
            "userId": user_id,
            "conversationId": conversation_id
        })
        
        print(f"第二个问题状态码: {response2.status_code}")
        print(f"第二个问题响应: {json.dumps(response2.json(), ensure_ascii=False, indent=2)}")
        print()
        
    except Exception as e:
        print(f"上下文聊天完成测试失败: {e}")
        print()

def test_root_endpoint():
    """测试根端点"""
    try:
        response = requests.get("http://localhost:8013/")
        print(f"根端点状态码: {response.status_code}")
        print(f"根端点响应: {response.json()}")
        print()
    except Exception as e:
        print(f"根端点测试失败: {e}")
        print()

if __name__ == "__main__":
    print("=== SQL Agent API 测试 ===\n")
    
    # 测试各个端点
    # test_root_endpoint()
    # test_health_check()
    # test_sql_query()
    # test_chat_completion()
    test_chat_completion_stream()
    # test_chat_completion_with_context()
    
    print("=== 测试完成 ===")